#pragma once

#include <iostream>
#include <thread>
#include <atomic>
#include <memory>
#include <vector>
#include "worker.hpp"
#include "functional.hpp"
#include "ring_queue.hpp"

template <typename Task, template <typename> class Queue>
class ThreadPoolImpl;
// using ThreadPool = ThreadPoolImpl<std::function<void()>, RingQueue>;
// using ThreadPool = ThreadPoolImpl<std::function<void()>, MPMCBoundedQueue>;
using ThreadPool = ThreadPoolImpl<Functional<void(), 128>, MPMCBoundedQueue>;
using ThreadPoolPtr = std::shared_ptr<ThreadPool>;

template <typename Task, template <typename> class Queue>
class ThreadPoolImpl
{
public:
    // ThreadPoolImpl(size_t thread_count = 2, size_t queue_size = 8192): _next_worker_id(0)
    ThreadPoolImpl(size_t thread_count = 2, size_t queue_size = 32768): _next_worker_id(0)
    {
        // std::cout << thread_count << std::endl;
        for (int i = 0; i < thread_count; i++)
        {
            _workers.emplace_back(std::make_shared<Worker<Task, Queue>>(queue_size));
        }

        for (int i = 0; i < thread_count; i++)
        {
            _workers[i]->start(i, _workers[(i + 1) % thread_count].get());
        }
    }

    template <typename Handler>
    void post(Handler &&handler)
    {
        const auto ok = tryPost(std::forward<Handler>(handler));
        if (!ok)
        {
            throw std::runtime_error("thread pool queue is full");
        }
    }

    template <typename Handler>
    bool tryPost(Handler &&handler)
    {
        return getWorker().post(std::forward<Handler>(handler));
    }

    ~ThreadPoolImpl()
    {
        for (auto worker_ptr : _workers)
        {
            worker_ptr->stop();
        }
    }

private:
    Worker<Task, Queue> &getWorker()
    {
        auto id = Worker<Task, Queue>::getWorkerIdForCurrentThread();

        if (id > _workers.size())
        {
            id = _next_worker_id.fetch_add(1, std::memory_order_relaxed) %
                 _workers.size();
        }

        return *_workers[id];
    }

    std::vector<std::shared_ptr<Worker<Task, Queue>>> _workers;
    std::atomic<size_t> _next_worker_id;
};
